'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Created on Jul 25, 2009

Version: 0.1

@author: bryan
'''

import JT

class ImageCanvas(JT.vtk.vtkRenderWindow):
    """
    This class was designed to be used for displaying fixed fluoro images.
    It might be useful to extend it so that it can be used as an image
    backround for DRR images.

    To use: set image filename path (SetFileName) and Render()
    """

    def __init__(self):
        """Initialize class
        """

        #Set up a render window for displaying the 2D fixed image
        self._ren = JT.vtk.vtkRenderer()
        self.AddRenderer(self._ren)
        self._iren = JT.vtk.vtkRenderWindowInteractor()
        self._iren.SetRenderWindow(self)

        #To make the itk and vtk coord sys align, the vtk camera is flipped
        self._cam = self._ren.GetActiveCamera()

        self._mapper = JT.vtk.vtkImageMapper()

        self._imageActor = JT.vtk.vtkActor2D()
        self._imageActor.SetMapper(self._mapper)

        self._ren.AddActor2D(self._imageActor)

        #Create an image Flipper, so that the image is displayed correctly
        self._flipperX = JT.vtk.vtkImageFlip()
        self._flipperX.SetFilteredAxis(0)

        self._flipperY = JT.vtk.vtkImageFlip()
        self._flipperY.SetFilteredAxis(1)

        self._verbose = False

    def SetImage(self,image):
        """
        Must pass VTK image data.  This method is called internally when
        the SetFileName function is called.

        Internally the image will be converted to VTK for display
        """
        self._image = image

        self._imageSize = []
        self._imageSize = image.GetDimensions()
        
        #for image to be displayed correctly, it has to be flipped about the X axis
        #Create a VTK Image Flipper
        flipperX = JT.vtk.vtkImageFlip()
        flipperY = JT.vtk.vtkImageFlip()
        flipperX.SetFilteredAxis(0)
        flipperY.SetFilteredAxis(1)
        flipperX.SetInput(image)
        flipperY.SetInput(flipperX.GetOutput())
        
        #self._imageMetric.SetFixedImage (self._itkImage )

        # Set size of the render window to be the same as the fixed image
        self.SetSize((self._imageSize[0:2]))

        self._mapper.SetInput(flipperY.GetOutput())

        # convert the VTK image to a numpy image array to perform some basic calculations
        # numpy support is a fast conversion, so this should be efficient
        # the converted numpy image is stored as a data member of this class for use
        # in other functions.
        pyImage = vtk2numpy(image)
        maxPixel = pyImage.max()

        self._mapper.SetColorLevel(maxPixel/2)
        self._mapper.SetColorWindow(maxPixel)

    def GetImage(self):
        """
        Return the current image (VTK format),
        python (numpy) image can be had by using self._
        """
        return self._image
        
    def SetFileName(self,filename):
        """
        Set the filename for the image to be displayed.

        It will be read using the appropriate VTK image file reader
        """
        if self._verbose:
            logging.debug("Image Canvas: opening file: %s" % filename)

        self._fileName =  filename
        
        image = genericImageReader(filename)
        self.SetImage(image)
        
    def Start(self):
        """
        Initiate interaction and render the image.
        """
        self.Render()
        self.SetWindowName('Image Canvas')
        self._iren.Start()

    def OverlayImages(self,image1,image2,outputFileName):
        """
        overlay two VTK images using contrasting colors (Red and Green) to visualize
        the differences in alignment.

        image1 will be Red, image2 will be Green
        """
        #Set up pipeline components
        lut1=JT.vtk.vtkLookupTable()
        lut2=JT.vtk.vtkLookupTable()

        colors1=JT.vtk.vtkImageMapToColors()
        colors2=JT.vtk.vtkImageMapToColors()

        imageMapper1=JT.vtk.vtkImageMapper()
        imageMapper2=JT.vtk.vtkImageMapper()

        blender=JT.vtk.vtkImageBlend()

        #First lookup table will map to red
        lut1.SetTableRange(0,3)
        lut1.SetHueRange(0,0)
        lut1.SetSaturationRange(0.0,1)
        lut1.SetValueRange(1,1)
        lut1.Build()

        #Sec(ond lookup table will map to green
        lut2.SetTableRange(0,3)
        lut2.SetHueRange(0.33,0.33)
        lut2.SetSaturationRange(0.0,1)
        lut2.SetValueRange(1,1)
        lut2.Build()

        #Make fixed image port 0 (RED) and Moving image port 1 (GREEN)
        colors1.SetInput(image1)
        colors1.SetLookupTable(lut1)

        colors2.SetInput(image2)
        colors2.SetLookupTable(lut2)

        imageMapper1.SetInput(colors1.GetOutput())
        imageMapper2.SetInput(colors2.GetOutput())

        imageActor1=JT.vtk.vtkActor2D()
        imageActor1.SetMapper(imageMapper1)

        imageActor2=JT.vtk.vtkActor2D()
        imageActor2.SetMapper(imageMapper2)

        blender.SetInput(0,colors1.GetOutput())
        blender.SetInput(1,colors2.GetOutput())

        #set the opacity of the ports so that both images are visible
        blender.SetOpacity(0,0.7)
        blender.SetOpacity(1,0.7)

        #Write the image to disk
        writer = JT.vtk.vtkPNGWriter()
        writer.SetInput(blender.GetOutput())
        writer.SetFileName(outputFileName)
        writer.Write()

    def CheckerImages(self,image1,image2,outputFileName):
        """
        Combine two images using a checkerboard pattern.  This is useful
        for visualizing the results of a registration.
        image1 and image2 should be VTK images
        outputFileName should be the path to save the result.
        """
        checker = JT.vtk.vtkImageCheckerboard()
        checker.SetNumberOfDivisions(8, 8, 1)
        
        checker.SetInput1(image1)
        checker.SetInput2(image2)
        
        writer = JT.vtk.vtkPNGWriter()
        writer.SetInput(checker.GetOutput())
        writer.SetFileName(outputFileName)
        writer.Write()
